import os
import numpy as np
from PIL import Image
from typing import Tuple, Optional

def preprocess_image(
    image_path: str,
    target_size: Tuple[int, int] = (224, 224),
    normalize: bool = True,
    mean: Optional[Tuple[float, float, float]] = (0.485, 0.456, 0.406),
    std: Optional[Tuple[float, float, float]] = (0.229, 0.224, 0.225)
) -> np.ndarray:
    """
    对图像进行预处理，使其符合模型输入要求
    :param image_path: 原始图像的本地路径
    :param target_size: 目标尺寸 (width, height)，默认(224, 224)
    :param normalize: 是否进行归一化，默认True
    :param mean: 归一化均值（与模型训练时保持一致），默认ImageNet均值
    :param std: 归一化标准差（与模型训练时保持一致），默认ImageNet标准差
    :return: 预处理后的图像数组，shape为(1, height, width, 3)
    :raises FileNotFoundError: 图像路径不存在时抛出
    :raises ValueError: 图像处理失败时抛出
    """
    # 检查文件是否存在
    if not os.path.exists(image_path):
        raise FileNotFoundError(f"图像文件不存在：{image_path}")

    try:
        # 1. 打开图像（使用with确保文件正确关闭）
        with Image.open(image_path) as img:
            # 2. 转换为RGB通道（处理RGBA/灰度图）
            if img.mode not in ('RGB', 'L', 'RGBA'):
                raise ValueError(f"不支持的图像模式：{img.mode}，仅支持RGB/L/RGBA")
            
            # 转换为RGB（灰度图会自动扩展为三通道，RGBA会丢弃Alpha通道）
            img_rgb = img.convert('RGB')

            # 3. 等比例缩放后居中裁剪（避免图像变形）
            # 计算缩放比例（按最短边缩放，确保图像能覆盖目标尺寸）
            img_width, img_height = img_rgb.size
            target_width, target_height = target_size

            # 计算缩放因子
            scale = max(target_width / img_width, target_height / img_height)
            new_width = int(img_width * scale)
            new_height = int(img_height * scale)

            # 缩放图像
            img_resized = img_rgb.resize((new_width, new_height), Image.Resampling.LANCZOS)  # 高质量缩放

            # 计算居中裁剪坐标
            left = (new_width - target_width) // 2
            top = (new_height - target_height) // 2
            right = left + target_width
            bottom = top + target_height

            # 裁剪到目标尺寸
            img_cropped = img_resized.crop((left, top, right, bottom))

            # 4. 转换为numpy数组（ dtype为float32，便于后续归一化 ）
            img_array = np.array(img_cropped, dtype=np.float32)

            # 5. 归一化（与模型训练时的预处理保持一致）
            if normalize:
                if mean is None or std is None:
                    raise ValueError("归一化时必须提供mean和std参数")
                img_array /= 255.0  # 先归一化到[0, 1]
                img_array = (img_array - mean) / std  # 标准化

            # 6. 增加批次维度（模型输入通常需要 (batch_size, height, width, channels)）
            img_array = np.expand_dims(img_array, axis=0)

            return img_array

    except Exception as e:
        raise ValueError(f"图像预处理失败：{str(e)}") from e
    


# -------------------------- 测试代码（修正版） --------------------------
if __name__ == "__main__":
    # 1. 定义测试参数（与 preprocess_image 函数的默认参数保持一致）
    TEST_IMAGE_PATH = "test_leaf.jpg"  # 测试图片路径（项目根目录）
    target_size = (224, 224)           # 目标尺寸
    use_normalize = True               # 是否归一化（对应函数的 normalize 参数）
    mean = (0.485, 0.456, 0.406)       # 归一化均值（对应函数的 mean 参数）
    std = (0.229, 0.224, 0.225)        # 归一化标准差（对应函数的 std 参数）

    # 2. 开始测试
    try:
        print("开始测试图像预处理...")
        
        # 调用预处理函数（显式传入参数，避免依赖默认值导致的变量未定义）
        processed_img = preprocess_image(
            image_path=TEST_IMAGE_PATH,
            target_size=target_size,
            normalize=use_normalize,
            mean=mean,
            std=std
        )
        
        # 3. 验证输出形状
        expected_shape = (1, target_size[0], target_size[1], 3)  # (1, 224, 224, 3)
        if processed_img.shape == expected_shape:
            print(f"✅ 形状验证通过：输出形状为 {processed_img.shape}（符合预期 {expected_shape}）")
        else:
            print(f"❌ 形状验证失败：实际形状 {processed_img.shape}，预期 {expected_shape}")
        
        # 4. 验证归一化（仅当 use_normalize 为 True 时）
        if use_normalize:
            min_val = processed_img.min()
            max_val = processed_img.max()
            # 归一化后的值应在 [-3, 3] 左右（因 ImageNet 均值/标准差）
            if -3 <= min_val <= 3 and -3 <= max_val <= 3:
                print(f"✅ 归一化验证通过：值范围 [{min_val:.2f}, {max_val:.2f}]（符合预期）")
            else:
                print(f"❌ 归一化验证异常：值范围 [{min_val:.2f}, {max_val:.2f}]（超出合理范围）")
        
        # 5. 保存预处理后的图像（用于可视化检查）
        if use_normalize and mean is not None and std is not None:
            # 还原归一化（将值从 [-3,3] 还原为 [0,255]，方便保存为图片）
            img_denorm = processed_img[0] * std + mean  # 第一步：还原标准化
            img_denorm = img_denorm * 255.0             # 第二步：还原到 [0,255]
            img_denorm = np.clip(img_denorm, 0, 255).astype(np.uint8)  # 确保值合法
            
            # 保存图像到项目根目录
            save_path = "preprocessed_test_leaf.jpg"
            Image.fromarray(img_denorm).save(save_path)
            print(f"📷 预处理后的图像已保存至：{save_path}（可打开查看效果）")
        
        print("\n🎉 测试完成！预处理函数工作正常")
    
    except FileNotFoundError as e:
        print(f"\n❌ 测试失败：{e}")
        print("👉 请检查：1. 图片路径是否正确 2. 图片文件是否存在")
    except ValueError as e:
        print(f"\n❌ 测试失败：{e}")
        print("👉 请检查：1. 图像格式是否支持（仅RGB/L/RGBA）2. 归一化参数是否正确")
    except Exception as e:
        print(f"\n❌ 测试时发生未知错误：{str(e)}")
        print("👉 请检查：1. 是否安装了所有依赖（pip install pillow numpy）2. 代码是否有拼写错误")